from dolfinx import mesh, plot
from dolfinx.fem import functionspace
from mpi4py import MPI
import pyvista as pv
import numpy as np


def plot_mesh(mesh: mesh.Mesh, values=None, output_file="mesh.png"):
    """
    给定DOLFINx网格，创建pyvista.UnstructuredGrid并保存为图像。

    参数:
        mesh: 要可视化的网格
        values: 每个网格单元的标记值列表
        output_file: 输出图像文件的路径
    """
    # 启动虚拟帧缓冲（用于无头服务器）
    pv.start_xvfb()

    # 创建pyvista绘图器实例
    plotter = pv.Plotter(off_screen=True)  # 设置为屏幕外渲染

    # 创建线性化网格
    V_linear = functionspace(mesh, ("Lagrange", 1))
    linear_grid = pv.UnstructuredGrid(*plot.vtk_mesh(V_linear))

    # 处理高阶网格
    if mesh.geometry.cmap.degree > 1:
        ugrid = pv.UnstructuredGrid(*plot.vtk_mesh(mesh))
        if values is not None:
            ugrid.cell_data["Marker"] = values
        plotter.add_mesh(ugrid, style="points", color="b", point_size=10)
        ugrid = ugrid.tessellate()
        plotter.add_mesh(ugrid, show_edges=False)
        plotter.add_mesh(linear_grid, style="wireframe", color="black")
    else:
        if values is not None:
            linear_grid.cell_data["Marker"] = values
        plotter.add_mesh(linear_grid, show_edges=True)

    # 设置视图
    plotter.show_axes()
    plotter.view_xy()

    # 保存图像而非显示
    plotter.screenshot(output_file)
    if MPI.COMM_WORLD.rank == 0:
        print(f"Mesh saved to {output_file}")


# 创建网格（仅在主进程中可视化）
msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(
        [1, 0])], n=[16, 16], cell_type=mesh.CellType.triangle
)

# 仅在主进程中调用可视化函数
if MPI.COMM_WORLD.rank == 0:
    plot_mesh(msh)
